from typing import List, Tuple, Dict, Optional, Set
from dataclasses import dataclass

@dataclass
class LayerInfo:
    """Information about a CNN layer."""
    layer_id: str
    layer_type: str
    kernel_size: Tuple[int, ...]
    stride: Tuple[int, ...]
    padding: Tuple[int, ...]
    dilation: Tuple[int, ...]
    input_shape: Tuple[int, ...]
    output_shape: Tuple[int, ...]

class RangeDeduction:
    """Implements Arbitrary Input Range (AIR) algorithm for CNN computation.
    
    This class provides methods to calculate minimal input ranges and intermediate
    layer output ranges for arbitrary output requirements.
    """
    
    def __init__(self):
        self.layers: Dict[str, LayerInfo] = {}
        self.layer_graph: Dict[str, List[str]] = {}  # layer_id -> list of next layer_ids
        self.intermediate_ranges: Dict[str, Tuple[int, int, int, int]] = {}
    
    def register_layer(self, layer_info: LayerInfo) -> None:
        """Register a new layer in the computation graph.
        
        Args:
            layer_info: Information about the layer
        """
        self.layers[layer_info.layer_id] = layer_info
    
    def add_layer_dependency(self, from_layer: str, to_layer: str) -> None:
        """Add a dependency between two layers.
        
        Args:
            from_layer: Source layer ID
            to_layer: Destination layer ID
        """
        if from_layer not in self.layer_graph:
            self.layer_graph[from_layer] = []
        self.layer_graph[from_layer].append(to_layer)
    
    def calculate_input_range(self, layer_id: str, output_range: Tuple[int, int, int, int]) \
            -> Tuple[int, int, int, int]:
        """Calculate the minimal input range for a given output range (AIR).
        
        Args:
            layer_id: Layer identifier
            output_range: (start_h, end_h, start_w, end_w) of required output
            
        Returns:
            Tuple[int, int, int, int]: Minimal input range (start_h, end_h, start_w, end_w)
        """
        layer = self.layers[layer_id]
        
        if layer.layer_type == 'conv2d' or layer.layer_type == 'pool2d':
            # For convolution layers with dilation
                        # 根据AIR算法公式(2)调整计算
            start_h = output_range[0] * layer.stride[0] - layer.padding[0]
            end_h = output_range[1] * layer.stride[0] + (layer.kernel_size[0] - 1) * layer.dilation[0] + 1 - layer.padding[0]
            
            # # 调试输出
            # print(f'\n=== 卷积层 {layer_id} 详细计算 ===')
            # print(f'公式参数: stride={layer.stride}  padding={layer.padding}  kernel={layer.kernel_size}  dilation={layer.dilation}')
            # print(f'输出范围: ({output_range[0]}, {output_range[1]}, {output_range[2]}, {output_range[3]})')
            
            # # 高度计算细节
            # print(f'\n高度方向计算:')
            # print(f'start_h = {output_range[0]}*{layer.stride[0]} - {layer.padding[0]} = {start_h}')
            # print(f'end_h = ({output_range[1]}-1)*{layer.stride[0]} + ({layer.kernel_size[0]}-1)*{layer.dilation[0]} + 1 - {layer.padding[0]}')
            # print(f'      = {output_range[1]-1}*{layer.stride[0]} + {layer.kernel_size[0]-1}*{layer.dilation[0]} + 1 - {layer.padding[0]}')
            # print(f'      = {(output_range[1]-1)*layer.stride[0]} + {(layer.kernel_size[0]-1)*layer.dilation[0]} + 1 - {layer.padding[0]}')
            # print(f'      = {end_h}')
            
            # 宽度计算
            start_w = output_range[2] * layer.stride[1] - layer.padding[1]
            end_w = (output_range[3]) * layer.stride[1] + (layer.kernel_size[1] - 1) * layer.dilation[1] + 1 - layer.padding[1]
            
            # print(f'\n宽度方向计算:')
            # print(f'start_w = {output_range[2]}*{layer.stride[1]} - {layer.padding[1]} = {start_w}')
            # print(f'end_w = ({output_range[3]}-1)*{layer.stride[1]} + ({layer.kernel_size[1]}-1)*{layer.dilation[1]} + 1 - {layer.padding[1]}')
            # print(f'      = {output_range[3]-1}*{layer.stride[1]} + {layer.kernel_size[1]-1}*{layer.dilation[1]} + 1 - {layer.padding[1]}')
            # print(f'      = {(output_range[3]-1)*layer.stride[1]} + {(layer.kernel_size[1]-1)*layer.dilation[1]} + 1 - {layer.padding[1]}')
            # print(f'      = {end_w}')
            
            # print(f'\n原始输入范围 (before clamping): ({start_h}, {end_h}, {start_w}, {end_w})')

            
        # elif layer.layer_type == 'pool2d':
        #     # For pooling layers
        #     start_h = output_range[0] * layer.stride[0]
        #     end_h = (output_range[1] - 1) * layer.stride[0] + layer.kernel_size[0]
        #     end_h = (output_range[1] - 1) * layer.stride[0] + (layer.kernel_size[0] - 1) * layer.dilation[0] + 1 - layer.padding[0]
        #     print(f'修正后end_h计算: ({output_range[1]}-1)*{layer.stride[0]} + {layer.kernel_size[0]} = {end_h}')
            
        #     start_w = output_range[2] * layer.stride[1]
        #     end_w = (output_range[3] - 1) * layer.stride[1] + layer.kernel_size[1]
        #     # print(f'修正后end_w计算: ({output_range[3]}-1)*{layer.stride[1]} + {layer.kernel_size[1]} = {end_w}')
            
        #     # print(f'\n=== 池化层 {layer_id} 详细计算 ===')
        #     # print(f'池化参数: stride={layer.stride}  kernel={layer.kernel_size}')
        #     # print(f'输出范围: ({output_range[0]}, {output_range[1]}, {output_range[2]}, {output_range[3]})')
        #     # print(f'计算后范围 (before clamping): ({start_h}, {end_h}, {start_w}, {end_w})')

            
        else:
            # For other layer types, use full input range
            start_h, end_h = 0, layer.input_shape[1] if len(layer.input_shape) > 1 else 0
            start_w, end_w = 0, layer.input_shape[3] if len(layer.input_shape) > 3 else 0
        
        # Clamp to valid input range and store
        # print(f'\n=== 输入范围限制 ===')
        # print(layer.input_shape)
        input_range = (
            max(0, start_h),
            min(layer.input_shape[1] if len(layer.input_shape) > 1 else 0, end_h),
            max(0, start_w),
            min(layer.input_shape[3] if len(layer.input_shape) > 3 else 0, end_w)
        )
        
        # Store intermediate range
        self.intermediate_ranges[layer_id] = input_range
        return input_range
    
    def AIR(self, input_layer_id: str, output_layer_id: str, out_range: tuple[int, int, int, int]) \
            -> tuple[dict[str, tuple[int, int, int, int]], dict[str, tuple[int, int, int, int]]]:
        """Perform Arbitrary Input Range (AIR) algorithm for a given output range."""
        # Step 1: 初始化中间范围集合
        intermediate_ranges = self.intermediate_ranges.copy()
        
        # Step 2: 检查输出层现有范围
        existing_output_range = intermediate_ranges.get(output_layer_id)
        if existing_output_range:
            # 合并范围 (start_h, end_h, start_w, end_w)
            merged_range = (
                min(out_range[0], existing_output_range[0]),
                max(out_range[1], existing_output_range[1]),
                min(out_range[2], existing_output_range[2]),
                max(out_range[3], existing_output_range[3])
            )
            if merged_range == existing_output_range:
                return {}, intermediate_ranges
            out_range = merged_range
        
        # Step 3: 输入层等于输出层
        if input_layer_id == output_layer_id:
            input_range = self.calculate_input_range(output_layer_id, out_range)
            return {input_layer_id: input_range}, intermediate_ranges
        
        # Step 4/5: 处理祖先层
        minimal_ranges = {}
        ancestors = self._get_all_ancestors(output_layer_id)
        
        for ancestor_id in ancestors:
            # 递归调用AIR算法
            ancestor_ranges, ancestor_intermediate = self.AIR(
                input_layer_id, 
                ancestor_id,
                self.calculate_input_range(output_layer_id, out_range)
            )
            
            # 合并中间结果
            intermediate_ranges.update(ancestor_intermediate)
            
            # 合并输入范围
            for layer_id, range_val in ancestor_ranges.items():
                if layer_id not in minimal_ranges:
                    minimal_ranges[layer_id] = range_val
                else:
                    # 范围合并逻辑
                    minimal_ranges[layer_id] = (
                        min(range_val[0], minimal_ranges[layer_id][0]),
                        max(range_val[1], minimal_ranges[layer_id][1]),
                        min(range_val[2], minimal_ranges[layer_id][2]),
                        max(range_val[3], minimal_ranges[layer_id][3])
                    )
        
        return minimal_ranges, intermediate_ranges
    
    def _get_all_ancestors(self, layer_id: str) -> List[str]:
        """获取指定层的所有祖先层（BFS遍历）"""
        visited = set()
        queue = [layer_id]
        
        while queue:
            current_id = queue.pop(0)
            if current_id not in visited:
                visited.add(current_id)
                # 获取直接前驱层
                predecessors = self.get_previous_layers(current_id)
                queue.extend(predecessors)
        
        return list(visited - {layer_id})

    def get_dependent_ranges(self, layer_id: str, target_range: Tuple[int, int, int, int]) \
            -> Dict[str, Tuple[int, int, int, int]]:
        """Get all dependent input ranges for previous layers using AIR algorithm.
        
        Args:
            layer_id: Current layer identifier
            target_range: Required output range for current layer
            
        Returns:
            Dict[str, Tuple[int, int, int, int]]: Map of layer_id to minimal input ranges
        """
        result = {}
        visited = set()
        
        def backtrack(current_id: str, current_range: Tuple[int, int, int, int]) -> None:
            if current_id in visited:
                return
                
            visited.add(current_id)
            
            # 实现AIR算法步骤2：范围合并
            existing_range = self.intermediate_ranges.get(current_id)
            if existing_range:
                # 合并当前范围与已有范围
                print(f'合并现有范围: 当前范围{current_range} vs 存储范围{existing_range}')
                merged_range = (
                    min(current_range[0], existing_range[0]),
                    max(current_range[1], existing_range[1]),
                    min(current_range[2], existing_range[2]),
                    max(current_range[3], existing_range[3])
                )
                print(f'合并后范围: {merged_range} (min_h,max_h,min_w,max_w)')
                if merged_range != existing_range:
                    current_range = merged_range
                    self.intermediate_ranges[current_id] = merged_range
            
            result[current_id] = current_range
            
            # 计算当前层输入范围
            input_range = self.calculate_input_range(current_id, current_range)
            
            # 反向传播处理依赖层（论文步骤5）
            for prev_layer in self.get_previous_layers(current_id):
                if prev_layer not in visited:
                    backtrack(prev_layer, input_range)
        
        backtrack(layer_id, target_range)
        return result
    
    def get_previous_layers(self, layer_id: str) -> List[str]:
        """Get list of layers that directly feed into the given layer.
        
        Args:
            layer_id: Layer identifier
            
        Returns:
            List[str]: List of previous layer IDs
        """
        return [prev_id for prev_id, next_layers in self.layer_graph.items()
                if layer_id in next_layers]
    
    def get_intermediate_ranges(self) -> Dict[str, Tuple[int, int, int, int]]:
        """Get all calculated intermediate layer ranges.
        
        Returns:
            Dict[str, Tuple[int, int, int, int]]: Map of layer_id to intermediate ranges
        """
        return self.intermediate_ranges.copy()
    
    def clear_intermediate_ranges(self) -> None:
        """Clear stored intermediate ranges."""
        self.intermediate_ranges.clear()